import numpy as np
import matplotlib.pyplot as plt

from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签


class DecisionTree:
    def __init__(self, dataset, args):
        self.dataset = dataset

        # Parameters
        self.n_classes = 3
        self.plot_colors = "ryb"
        self.plot_step = 0.02

        # Load data
        self.iris = load_iris()

    def run(self):
        for pairidx, pair in enumerate([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]]):
            # We only take the two corresponding features
            X = self.iris.data[:, pair]
            y = self.iris.target

            # Train
            clf = DecisionTreeClassifier().fit(X, y)

            # Plot the decision boundary
            plt.subplot(2, 3, pairidx + 1)

            x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
            y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
            xx, yy = np.meshgrid(np.arange(x_min, x_max, self.plot_step),
                                 np.arange(y_min, y_max, self.plot_step))
            plt.tight_layout(h_pad=0.5, w_pad=0.5, pad=2.5)

            Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])
            Z = Z.reshape(xx.shape)
            cs = plt.contourf(xx, yy, Z, cmap=plt.cm.RdYlBu)

            plt.xlabel(self.iris.feature_names[pair[0]])
            plt.ylabel(self.iris.feature_names[pair[1]])

            # Plot the training points
            for i, color in zip(range(self.n_classes), self.plot_colors):
                idx = np.where(y == i)
                plt.scatter(X[idx, 0], X[idx, 1], c=color, label=self.iris.target_names[i],
                            cmap=plt.cm.RdYlBu, edgecolor='black', s=15)

        plt.suptitle("应用不同特征集构建的决策树")
        plt.legend(loc='lower right', borderpad=0, handletextpad=0)
        plt.axis("tight")

        # plt.figure()
        # clf = DecisionTreeClassifier().fit(self.iris.data, self.iris.target)
        # plot_tree(clf, filled=True)
        plt.savefig('../../result/dt_%s.png' % self.dataset)
